from keras.layers import LSTM, Input, Concatenate, Reshape, Lambda, RepeatVector
from keras.models import Model
from GraphAttention import GraphAttention as GAT

n_feature = 7
input_timeStep = 3
encode_size = 64
def generateModel(carNum):
    # 邻接矩阵数据
    A = Input(shape=(carNum, carNum))
    # 每辆车数据
    allIntput = []
    for i in range(carNum):
        allIntput.append(Input(shape=(input_timeStep, n_feature)))
    # LSTM编码器
    allLstmEncoder = []
    for i in range(carNum):
        lstmEncoder = LSTM(encode_size, activation='relu', return_sequences=False)(allIntput[i])
        allLstmEncoder.append(lstmEncoder)
    # GAT
    X = Concatenate(axis=1)(allLstmEncoder)
    X = Reshape((input_timeStep, encode_size))(X)
    model = GAT(encode_size)([X, A])
    return Model(inputs=allIntput+[A], outputs=model)

testModel = generateModel(3)
print(testModel.summary())